#include <mm.h>
#include <fmt.h>

static void page_test();

void page_init()
{   
    printf("++ memory init\n");
    uintptr_t _heap_start_aligned = _align_page(HEAP_START);
    uint32_t num_reserved_page = LENGTH_RAM / (PAGE_SIZE * PAGE_SIZE);
    _num_pages = (HEAP_SIZE - (_heap_start_aligned - HEAP_START)) / PAGE_SIZE - num_reserved_page;
    printf("HEAP_START: 0x%lx, HEAP_SIZE: 0x%lx\n"
           "Alloc pages: [0x%x]\n",
           HEAP_START, HEAP_SIZE, _num_pages);

    struct Page *page = (struct Page *)HEAP_START;
    for (int i = 0; i < _num_pages; i ++) {
        _clear(page);
        page++;
    }

    _alloc_start = _heap_start_aligned + num_reserved_page * PAGE_SIZE;
    _alloc_end = _alloc_start + (PAGE_SIZE * _num_pages);

    printf("Alloc memory: [0x%lx - 0x%lx]\n", _alloc_start, _alloc_end);

    page_test();
    printf("page test success.\n");
}

void *page_alloc(int npages) {
    bool found = false;
    struct Page *page = (struct Page *)HEAP_START;
    for (int i = 0; i + npages <= _num_pages; i++) {
        if (_is_free(page + i)) {
            found = true;
            for (int j = i + 1; j < i  + npages; j ++) {
                if (!_is_free(page + j)) {
                    found = false;
                    break;
                }
            }

            if (found) {
                struct Page *p = page + i;
                for (int j = i; j < i + npages; j ++) {
                    _set_flag(p + j, PAGE_USED);
                }
                _set_flag(p + i + npages - 1, PAGE_LAST);
                return (void *)(_alloc_start + i * PAGE_SIZE);
            }
        }
    }
    return NULL;
}

void page_free(void *p) {
    if (!p || ((uintptr_t)p < _alloc_start && (uintptr_t)p >= _alloc_end)) {
        return;
    }
    struct Page *page = (struct Page *)HEAP_START;
    int page_num = ((uintptr_t)p - _alloc_start)/ PAGE_SIZE;

    int i = page_num;
    while (!_is_last(page + i)) {
        _clear(page + i);
        i ++;
    }
    _clear(page + i);
}

static int alloc_num_test() {
    struct Page *p = (struct Page *)HEAP_START;
    int cnt = 0;
    for (int i = 0; i < _num_pages; i ++) {
        if (!_is_free(p + i)) {
            printf("page [%d]: %d\n", i, (p + i)->flags);
            cnt ++;
        }
    }
    return cnt;
}

static void page_test()
{
    void *p = page_alloc(2);
    page_free(p);
    assert(alloc_num_test(p) == 0);

	void *p2 = page_alloc(2);
    page_free(p2);
    assert(alloc_num_test(p) == 0);

    // void *p3 = page_alloc(_num_pages);
	// assert(p3 != NULL);
    // assert(page_alloc(1) == NULL);
    // page_free(p3);
    // assert(alloc_num_test(p) == 0);
}